import argparse
from collections import defaultdict, deque

import cv2
import numpy as np
from ultralytics import YOLO

import supervision as sv
#python 1.py --source_video_path "D:\BaiduNetdiskDownload\MW4X100ImF\MW4X100ImF-1.MOV" --target_video_path "D:\yolo\ultralytics\target_output.MOV"

#需要根据原视频的四个点确认，原始应该为左上右上，右下左下> 但是泳池方向朝左，需要转化为右上右下、左下左上
SOURCE = np.array([ [1580,349], [1899,968], [486,1006],[680,373]])
#将四个点转换为宽21米长15米的矩形
TARGET_WIDTH = 21
TARGET_HEIGHT = 15

TARGET = np.array(
    [
        [0, 0],
        [TARGET_WIDTH - 1, 0],
        [TARGET_WIDTH - 1, TARGET_HEIGHT - 1],
        [0, TARGET_HEIGHT - 1],
    ]
)



class ViewTransformer:
    def __init__(self, source: np.ndarray, target: np.ndarray) -> None:
        source = source.astype(np.float32)
        target = target.astype(np.float32)
        self.m = cv2.getPerspectiveTransform(source, target)

    def transform_points(self, points: np.ndarray) -> np.ndarray:
        """
        视图转换器类，用于将点集从一个视图转换到另一个视图。
        该类通过计算两个视图之间的透视变换矩阵，然后应用这个变换矩阵来转换点集。使得2维空间的坐标转换到3维中。
        参数:
        source: np.ndarray, 源视图中的四个角点，形成一个3x4的矩阵。
        target: np.ndarray, 目标视图中的四个角点，形成一个3x4的矩阵。
        """
        if points.size == 0:
            return points

        reshaped_points = points.reshape(-1, 1, 2).astype(np.float32)
        transformed_points = cv2.perspectiveTransform(reshaped_points, self.m)
        return transformed_points.reshape(-1, 2)


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Swimming speed estimation"
    )
    parser.add_argument(
        "--source_video_path",
        required=True,
        help="Path to the source video file",
        type=str,
    )
    parser.add_argument(
        "--target_video_path",
        required=True,
        help="Path to the target video file (output)",
        type=str,
    )
    parser.add_argument(
        "--confidence_threshold",
        default=0.3,
        help="Confidence threshold for the model",
        type=float,
    )
    parser.add_argument(
        "--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
    )

    return parser.parse_args()


if __name__ == "__main__":
    args = parse_arguments()

    video_info = sv.VideoInfo.from_video_path(video_path=args.source_video_path)
    model = YOLO(r"D:\yolo\ultralytics\runs\detect\train3\weights\游泳522.pt")

    byte_track = sv.ByteTrack(
        frame_rate=video_info.fps, track_thresh=args.confidence_threshold
    )

    # 计算最优线条厚度，用于视频中的线条标注，确保线条粗细适中且清晰可见
    thickness = sv.calculate_optimal_line_thickness(
        resolution_wh=video_info.resolution_wh
    )
    # 计算最优文本尺寸，用于视频中的文本标注，确保文本大小适中且清晰可见
    text_scale = sv.calculate_optimal_text_scale(resolution_wh=video_info.resolution_wh)
    # 初始化边界框标注器，用于在视频中绘制边界框，线条厚度由thickness参数决定
    bounding_box_annotator = sv.BoundingBoxAnnotator(thickness=thickness)
    # 初始化标签标注器，用于在视频中绘制文本标签，文本大小由text_scale决定，线条厚度由thickness决定
    label_annotator = sv.LabelAnnotator(
        text_scale=text_scale,
        # text_thickness=thickness,
        text_thickness=2,
        text_position=sv.Position.BOTTOM_CENTER,
    )
    # 初始化轨迹标注器，用于在视频中绘制轨迹，线条厚度由thickness决定，轨迹长度由视频帧率乘以2决定
    trace_annotator = sv.TraceAnnotator(
        thickness=thickness,
        trace_length=video_info.fps * 2,
        position=sv.Position.BOTTOM_CENTER,
    )

    # 从源视频路径获取视频帧生成器
    frame_generator = sv.get_video_frames_generator(source_path=args.source_video_path)

    # 初始化多边形区域对象，用于后续筛选检测结果
    polygon_zone = sv.PolygonZone(polygon=SOURCE)
    # 初始化视图转换器，用于将检测结果从源视图转换到目标视图
    view_transformer = ViewTransformer(source=SOURCE, target=TARGET)

    # 使用defaultdict和deque维护每个跟踪对象的纵坐标历史记录，以计算速度
    coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))

    # 使用VideoSink打开目标视频文件，准备写入处理后的帧
    with sv.VideoSink(args.target_video_path, video_info) as sink:
        # 遍历生成器提供的每一帧
        for frame in frame_generator:
            # 使用模型对当前帧进行检测，获取结果
            result = model(frame)[0]
            # 从模型结果中创建检测对象
            detections = sv.Detections.from_ultralytics(result)
            # 筛选出置信度高于指定阈值的检测
            detections = detections[detections.confidence > args.confidence_threshold]
            # 过滤出位于多边形区域内的检测
            detections = detections[polygon_zone.trigger(detections)]
            # 对检测结果进行非极大值抑制，减少重叠区域的检测
            detections = detections.with_nms(threshold=args.iou_threshold)
            # 使用ByteTrack更新跟踪信息，并获取更新后的检测结果
            detections = byte_track.update_with_detections(detections=detections)

            # 获取每个检测对象在底部中心位置的坐标
            points = detections.get_anchors_coordinates(
                anchor=sv.Position.BOTTOM_CENTER
            )
            # 将坐标从源视图转换到目标视图
            points = view_transformer.transform_points(points=points).astype(int)

            # 更新每个跟踪对象的坐标历史记录
            for tracker_id, [_, y] in zip(detections.tracker_id, points):
                coordinates[tracker_id].append(y)

            # 根据坐标历史记录计算速度，并为每个检测对象生成标签
            labels = []
            for tracker_id in detections.tracker_id:
                # 如果坐标记录不足半帧，只显示跟踪ID
                if len(coordinates[tracker_id]) < video_info.fps / 2:
                    labels.append(f"#{tracker_id}")
                else:
                    # 计算速度，单位为m/s
                    coordinate_start = coordinates[tracker_id][-1]
                    coordinate_end = coordinates[tracker_id][0]
                    distance = abs(coordinate_start - coordinate_end)
                    time = len(coordinates[tracker_id]) / video_info.fps
                    speed = distance / time
                    speed_str = f"{speed:.2f}"#速度保留两位小数不四舍五入
                    speed_str = f"{round(speed, 2):.2f}"#保留两位小数并且最后一位四舍五入

                    labels.append(f"#{tracker_id} {speed_str} m/s")

            # 依次使用不同注释器对帧进行注释
            annotated_frame = frame.copy()
            annotated_frame = trace_annotator.annotate(
                scene=annotated_frame, detections=detections
            )
            annotated_frame = bounding_box_annotator.annotate(
                scene=annotated_frame, detections=detections
            )
            annotated_frame = label_annotator.annotate(
                scene=annotated_frame, detections=detections, labels=labels
            )

            # 将处理后的帧写入目标视频文件
            sink.write_frame(annotated_frame)
            # 显示处理后的帧，等待退出指令
            cv2.imshow("frame", annotated_frame)
            if cv2.waitKey(1) & 0xFF == ord("q"):
                break
        # 关闭所有OpenCV窗口
        cv2.destroyAllWindows()
